Walkthrough

Adds a new PassConfig option to allow unswitching when else-branches have side effects, introduces let-binding-aware variable usage checks and side-effect detection, and extends the loop-unswitching pass to track per-thread (threadIdx) variables and honor the new config when deciding to hoist conditionals.

Changes

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related issues

• [BUG] Loop unswitching pass causes buggy codegen #1804 — Related: this PR’s per-thread-aware checks and conservative default for non-trivial else hoisting address the safety concerns described in that issue.

Possibly related PRs

• [Feature] Implement LoopUnswitching Pass #1747 — Foundational LoopUnswitching pass and initial pass-config integration that this change extends.
• [Feature] Enhance Loop Unswitching with Let Binding and Condition Handling #1766 — Related enhancements around Let-bound variable handling and If-condition hoisting used by this PR.
• [BugFix] Missing Recursive Loop Var Checking in Loop Unswitching #1801 — Adds/addresses recursive let-binding variable usage checking similar to UsesVarsThroughLetBindings.

Suggested reviewers

• chengyupku

Poem

🐰 I hop where loops and let-bindings play,
ThreadIdx warns me: "Don't unswitch today."
A config nudge, a careful, clever tweak,
Now hoists are safe — the rabbit gives a squeak! 🥕

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

• 📝 Generate docstrings

src/transform/loop_unswitching.cc (3)

526-526: thread_idx_vars_in_scope_ should be private.

This is an internal implementation detail of the mutator. Being public allows accidental access from outside the class.

Move to private section

 class LoopUnswitcher : public StmtExprMutator {
 public:
   explicit LoopUnswitcher(bool allow_non_trivial_else)
       : allow_non_trivial_else_(allow_non_trivial_else) {}
 
-  std::unordered_set<const VarNode *> thread_idx_vars_in_scope_;
-
   Stmt VisitStmt_(const ForNode *op) final {
   ...
 
 private:
   bool allow_non_trivial_else_{false};
+  std::unordered_set<const VarNode *> thread_idx_vars_in_scope_;
 };

---

528-629: Consider an RAII guard for thread_idx_vars_in_scope_ cleanup to avoid duplication across four early-return sites.

The pushed_thread_idx / erase pair is repeated at lines 559, 579, 601, and 625. A scope guard would consolidate this and prevent omission in future edits.

Example with a simple scope guard

   Stmt VisitStmt_(const ForNode *op) final {
     bool pushed_thread_idx = false;
     if (op->thread_binding.defined()) {
       String thread_tag = op->thread_binding.value()->thread_tag;
       if (thread_tag == "threadIdx.x" || thread_tag == "threadIdx.y" ||
           thread_tag == "threadIdx.z") {
         thread_idx_vars_in_scope_.insert(op->loop_var.get());
         pushed_thread_idx = true;
       }
     }
+    // RAII cleanup for thread_idx tracking
+    auto cleanup = [&]() {
+      if (pushed_thread_idx) {
+        thread_idx_vars_in_scope_.erase(op->loop_var.get());
+      }
+    };
+    struct ScopeGuard {
+      std::function<void()> fn;
+      ~ScopeGuard() { fn(); }
+    } guard{cleanup};

Then remove the four individual if (pushed_thread_idx) { ... erase ... } blocks and the manual cleanup before each return.

---

296-346: AttrStmtNode is not handled — may skip valid optimizations.

IsSideEffectFreeStmt returns false for any AttrStmtNode. Since HoistableIfFinder will descend into AttrStmt bodies and find loop-invariant ifs, but IfBranchReplacer preserves the AttrStmt wrapper when replacing the matched if with Evaluate(0), the resulting else-branch will contain AttrStmt(..., Evaluate(0)) which fails the side-effect-free check and blocks unswitching.

While the conservative approach is intentional, AttrStmt is a common annotation in TVM TIR for thread_extent, storage_scope, and async regions. Other transform passes in the codebase routinely check AttrStmtNode; handling it would enable optimizations in those patterns.

Consider adding:

   if (const auto *op = stmt.as<ForNode>()) {
     if (SideEffect(op->min) > CallEffectKind::kReadState ||
         SideEffect(op->extent) > CallEffectKind::kReadState) {
       return false;
     }
     return IsSideEffectFreeStmt(op->body);
   }
 
+  if (const auto *op = stmt.as<AttrStmtNode>()) {
+    if (SideEffect(op->value) > CallEffectKind::kReadState) {
+      return false;
+    }
+    return IsSideEffectFreeStmt(op->body);
+  }
+
   // Conservatively treat all other statements as side-effecting.
   return false;

• X
• Mastodon
• Reddit
• LinkedIn